package encrypt

import (
	"bytes"
	"crypto/aes"
	"crypto/cipher"
	"crypto/rand"
	"encoding/base64"
	"encoding/hex"
	"io"
)

type N新ase struct {
	ASE密钥        []byte
	V偏移量         []byte
	V加密块         cipher.Block
	V请求方法, V填充方式 string //支持CBC CFB ECB
	V加密结果方法      string //支持base64 hex
}

const (
	V请求方法_CBC = "CBC"
	V请求方法_CFB = "CFB"
	V请求方法_ECB = "ECB"

	V加解密方式_base64 = "base64"
	V加解密方式_hex    = "hex"

	V填充方式_pkcs5Padding = "pkcs5Padding"
	V填充方式_pkcs7Padding = "pkcs7Padding"
)

func F新建Ase配置(args ...string) S加解密接口 {
	conf := &N新ase{}
	lenArgs := len(args)
	switch lenArgs {
	case 1, 2, 3, 4, 5:
		conf.ASE密钥 = F处理密钥格式(args[0])
		if lenArgs >= 2 {
			conf.V请求方法 = args[1]
		}
		if lenArgs >= 3 {
			conf.V加密结果方法 = args[2]
		}
		if lenArgs >= 4 {
			conf.V填充方式 = args[3]
		}
		if lenArgs >= 5 {
			conf.V偏移量 = []byte(args[4])
		}
	default: //如果为0或者其它的就用默认的
		conf.ASE密钥 = []byte(默认ASE密钥)
	}
	conf.V加密块, _ = aes.NewCipher(conf.ASE密钥)
	return conf
}

type I设置接口 interface {
	apply(*N新ase)
}

type F设置方法 struct {
	f func(*N新ase)
}

func F新建ase(opts ...I设置接口) S加解密接口 {
	cc := &N新ase{ASE密钥: []byte(默认ASE密钥)}
	//循环调用opts
	for _, opt := range opts {
		opt.apply(cc)
	}
	cc.V加密块, _ = aes.NewCipher(cc.ASE密钥)
	return cc
}

func F新建ase设置方法(f func(*N新ase)) *F设置方法 {
	return &F设置方法{
		f: f,
	}
}

func F设置密钥(s string) *F设置方法 {
	return F新建ase设置方法(func(o *N新ase) {
		o.ASE密钥 = []byte(s)
	})
}
func F设置偏移量(s string) *F设置方法 {
	return F新建ase设置方法(func(o *N新ase) {
		o.V偏移量 = []byte(s)
	})
}
func F设置加密请求方法(s string) *F设置方法 {
	return F新建ase设置方法(func(o *N新ase) {
		o.V请求方法 = s
	})
}
func F设置加密填充方式(s string) *F设置方法 {
	return F新建ase设置方法(func(o *N新ase) {
		o.V填充方式 = s
	})
}
func F设置加密加解密方式(s string) *F设置方法 {
	return F新建ase设置方法(func(o *N新ase) {
		o.V加密结果方法 = s
	})
}

func (a *N新ase) F加密信息(body string) (string, error) {
	var by []byte
	var err error

	switch a.V请求方法 {
	case V请求方法_CBC:
		by, err = a.AesEncryptCBC([]byte(body))
	case V请求方法_CFB:
		by, err = a.AesEncryptCFB([]byte(body))
	case V请求方法_ECB:
		by, err = a.AesEncryptECB([]byte(body))
	default:
		by, err = a.AesEncryptCBC([]byte(body))
	}
	if err != nil {
		return "", err
	}

	switch a.V加密结果方法 {
	case V加解密方式_base64:
		return base64.StdEncoding.EncodeToString(by), nil
	case V加解密方式_hex:
		return hex.EncodeToString(by), nil
	default:
		return base64.StdEncoding.EncodeToString(by), nil
	}

}
func (a *N新ase) F解密信息(body string) (string, error) {
	var by []byte
	var err error

	switch a.V加密结果方法 {
	case V加解密方式_base64:
		by, err = base64.StdEncoding.DecodeString(body)
	case V加解密方式_hex:
		by, err = hex.DecodeString(body)
	default:
		by, err = base64.StdEncoding.DecodeString(body)
	}
	if err != nil {
		return "", err
	}
	switch a.V请求方法 {
	case V请求方法_CBC:
		by, err = a.AesDecryptCBC(by)
	case V请求方法_CFB:
		by, err = a.AesDecryptCFB(by)
	case V请求方法_ECB:
		by, err = a.AesDecryptECB(by)
	default:
		by, err = a.AesDecryptCBC(by)
	}
	if err != nil {
		return "", err
	}
	return string(by), nil
}

func (a *N新ase) F获取加密补全方法() func(ciphertext []byte, blockSize int) []byte {
	switch a.V填充方式 {
	case V填充方式_pkcs7Padding:
		return PKCS7Padding
	case V填充方式_pkcs5Padding:
		return pkcs5Padding
	default:
		return pkcs5Padding
	}
}
func (a *N新ase) F获取解密补全方法() func(origData []byte) []byte {
	switch a.V填充方式 {
	case V填充方式_pkcs7Padding:
		return PKCS7UnPadding
	case V填充方式_pkcs5Padding:
		return pkcs5UnPadding
	default:
		return pkcs5UnPadding
	}
}

func (a *N新ase) F获取偏移量() (iv []byte) {
	if len(a.V偏移量) == 0 {
		return a.ASE密钥[:a.V加密块.BlockSize()]
	} else {
		return a.V偏移量
	}
}

// =================== CBC ======================
func (a *N新ase) AesEncryptCBC(origData []byte) (encrypted []byte, err error) {
	// 分组秘钥
	// NewCipher该函数限制了输入k的长度必须为16, 24或者32

	blockSize := a.V加密块.BlockSize() // 获取秘钥块的长度
	origData = a.F获取加密补全方法()(origData, blockSize)
	blockMode := cipher.NewCBCEncrypter(a.V加密块, a.F获取偏移量()) // 加密模式
	encrypted = make([]byte, len(origData))                 // 创建数组
	blockMode.CryptBlocks(encrypted, origData)              // 加密
	return
}

func (a *N新ase) AesDecryptCBC(encrypted []byte) (decrypted []byte, err error) {
	// 获取秘钥块的长度
	blockMode := cipher.NewCBCDecrypter(a.V加密块, a.F获取偏移量()) // 加密模式
	decrypted = make([]byte, len(encrypted))                // 创建数组
	blockMode.CryptBlocks(decrypted, encrypted)             // 解密
	decrypted = a.F获取解密补全方法()(decrypted)                    // 去除补全码
	return
}
func pkcs5Padding(ciphertext []byte, blockSize int) []byte {
	padding := blockSize - len(ciphertext)%blockSize
	padtext := bytes.Repeat([]byte{byte(padding)}, padding)
	return append(ciphertext, padtext...)
}
func pkcs5UnPadding(origData []byte) []byte {
	length := len(origData)
	unpadding := int(origData[length-1])
	return origData[:(length - unpadding)]
}

// =================== ECB ======================
func (a *N新ase) AesEncryptECB(origData []byte) (encrypted []byte, err error) {

	length := (len(origData) + aes.BlockSize) / aes.BlockSize
	plain := make([]byte, length*aes.BlockSize)
	copy(plain, origData)
	pad := byte(len(plain) - len(origData))
	for i := len(origData); i < len(plain); i++ {
		plain[i] = pad
	}
	encrypted = make([]byte, len(plain))
	// 分组分块加密
	for bs, be := 0, a.V加密块.BlockSize(); bs <= len(origData); bs, be = bs+a.V加密块.BlockSize(), be+a.V加密块.BlockSize() {
		a.V加密块.Encrypt(encrypted[bs:be], plain[bs:be])
	}

	return
}
func (a *N新ase) AesDecryptECB(encrypted []byte) (decrypted []byte, err error) {

	decrypted = make([]byte, len(encrypted))
	for bs, be := 0, a.V加密块.BlockSize(); bs < len(encrypted); bs, be = bs+a.V加密块.BlockSize(), be+a.V加密块.BlockSize() {
		a.V加密块.Decrypt(decrypted[bs:be], encrypted[bs:be])
	}
	trim := 0
	if len(decrypted) > 0 {
		trim = len(decrypted) - int(decrypted[len(decrypted)-1])
	}
	decrypted = decrypted[:trim]
	return
}
func F处理密钥(keystr string) (genKey []byte) {
	l := len(keystr)
	if l == 16 || l == 24 || l == 32 {
		//如果位数是16 24 32位，直接使用
	} else {
		if l < 16 {
			keystr = keystr + 默认ASE密钥
		}
		l = 16
	}
	key := []byte(keystr)
	genKey = make([]byte, l)
	copy(genKey, key)
	for i := l; i < len(key); {
		for j := 0; j < l && i < len(key); j, i = j+1, i+1 {
			genKey[j] ^= key[i]
		}
	}
	return genKey
}

// =================== CFB ======================
func (a *N新ase) AesEncryptCFB(origData []byte) (encrypted []byte, err error) {
	encrypted = make([]byte, aes.BlockSize+len(origData))
	//iv := a.F获取偏移量()
	iv := encrypted[:aes.BlockSize]
	if _, err := io.ReadFull(rand.Reader, iv); err != nil {
		panic(err)
	}
	stream := cipher.NewCFBEncrypter(a.V加密块, iv)
	stream.XORKeyStream(encrypted[aes.BlockSize:], origData)
	return
}
func (a *N新ase) AesDecryptCFB(encrypted []byte) (decrypted []byte, err error) {

	iv := encrypted[:aes.BlockSize]
	encrypted = encrypted[aes.BlockSize:]
	//iv = a.F获取偏移量()
	stream := cipher.NewCFBDecrypter(a.V加密块, iv)
	stream.XORKeyStream(encrypted, encrypted)

	return encrypted, nil
}
